from TextComparator import BertScoreComparator, RougeComparator
import argparse
import json
from util import JsonUtil

class ArgumentParser:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description='Compare two texts.')
        self.parser.add_argument('--compare_path', type=str, required=True,
                                 help='the path of the text to compare')
        self.parser.add_argument('--comparator', type=str, choices=['bert_score', 'rouge_l'], required=True,
                                 help='the type of the text comparator')
        # 请把最终的结果文件存入result文件夹
        self.parser.add_argument('--output_path', type=str, required=True,
                                 help='the path of the output file')

    def parse_args(self):
        args = self.parser.parse_args()

        # 根据参数选择文本比较器
        if args.comparator == 'bert_score':
            comparator = BertScoreComparator()
        elif args.comparator == 'rouge_l':
            comparator = RougeComparator()

        return args.compare_path, comparator, args.output_path


if __name__ == '__main__':
    arg_parser = ArgumentParser()
    compare_path, comparator, output_path = arg_parser.parse_args()
    json_util = JsonUtil(output_path)
    labels = ['P', 'R', 'F1']

    # 打开json文件，读取json文件中的每一个data对象，取出generated_origin和generated_llm字段,使用comparator进行比较
    with open(compare_path, 'r', encoding='utf-8') as f:
        for line in f:
            data = json.loads(line)
            ref = data['generated_origin']
            hyp = data['generated_llm']

            # 检查ref和hyp是否为null
            if ref is None or hyp is None:
                continue

            score = comparator.compare(hyp, ref)

            # 如果是bert_score，则进行tensor格式转变
            if comparator.__class__.__name__ == 'BertScoreComparator':
                score_result = [f'{label}: {score.item():.4f}' for label, score in zip(labels, score)]
                # 将score作为新的字段加入到data中, 在此之前需要将score由Tensor类型转化为字符串类型
                data['score'] = ', '.join(score_result)

            else:
                # score此时是一个字典，包含了rouge-l的P, R, F1三个值，将score作为新字段加入到data中
                data['score'] = ', '.join([f'{label}: {score}' for label, score in score.items()])

            # 将data写入到output_path文件中
            json_util.json_write(data)
            
